#!/usr/bin/env python3

import rclpy
from rclpy.node import Node
from geometry_msgs.msg import Twist
from nav_msgs.msg import Odometry
from std_msgs.msg import Float64


class SimpleZVelocityController(Node):
    """简单的z轴速度控制器 - 固定速度0.2m/s"""
    
    def __init__(self):
        super().__init__('simple_z_velocity_controller')
        
        # 声明参数
        self.declare_parameter('z_velocity', 0.2)
        self.declare_parameter('altitude_tolerance', 0.05)
        
        # 获取参数
        self.z_velocity = self.get_parameter('z_velocity').get_parameter_value().double_value
        self.altitude_tolerance = self.get_parameter('altitude_tolerance').get_parameter_value().double_value
        
        # 状态变量
        self.target_altitude = 0.0
        self.current_altitude = 0.0
        self.altitude_reached = True
        self.pose_received = False
        self.target_set = False  # 标记是否设置过目标高度
        
        # 订阅器
        self.target_altitude_sub = self.create_subscription(
            Float64,
            '/target_altitude',
            self.target_altitude_callback,
            10
        )
        
        # 从里程计获取当前高度
        self.odom_sub = self.create_subscription(
            Odometry,
            '/red_standard_robot1/odometry',
            self.odometry_callback,
            10
        )
        
        # 发布器 - 直接发布到cmd_vel
        self.cmd_vel_pub = self.create_publisher(
            Twist,
            '/red_standard_robot1/cmd_vel',
            10
        )
        
        # 控制定时器 - 50Hz
        self.control_timer = self.create_timer(0.02, self.control_loop)
        
        # 调试定时器 - 每1秒
        self.debug_timer = self.create_timer(1.0, self.debug_output)
        
        self.get_logger().info('简单z轴速度控制器已启动')
        self.get_logger().info(f'固定z轴速度: ±{self.z_velocity} m/s')
        self.get_logger().info(f'高度容差: ±{self.altitude_tolerance} m')
        
    def target_altitude_callback(self, msg):
        """目标高度回调"""
        new_target = msg.data
        if abs(new_target - self.target_altitude) > 0.01:  # 1cm的变化阈值
            self.target_altitude = new_target
            self.altitude_reached = False
            self.target_set = True  # 标记已设置目标
            self.get_logger().info(f'新目标高度: {self.target_altitude:.2f}m')
            
    def odometry_callback(self, msg):
        """里程计回调"""
        self.current_altitude = msg.pose.pose.position.z
        self.pose_received = True
        
    def control_loop(self):
        """主控制循环"""
        # 创建速度命令
        cmd_vel = Twist()
        cmd_vel.linear.z = 0.0  # 默认不移动
        
        if not self.pose_received:
            # 如果没有位置数据，停止z轴运动
            self.cmd_vel_pub.publish(cmd_vel)
            return
            
        if not self.target_set:
            # 如果还没设置过目标高度，不发布任何z轴速度
            self.cmd_vel_pub.publish(cmd_vel)
            return
            
        # 计算高度误差
        error = self.target_altitude - self.current_altitude
        
        # 检查是否到达目标
        if abs(error) <= self.altitude_tolerance:
            # 已到达目标，停止z轴运动
            cmd_vel.linear.z = 0.0
            if not self.altitude_reached:
                self.altitude_reached = True
                self.get_logger().info(f'高度已到达目标: {self.current_altitude:.3f}m')
        else:
            # 没有到达目标，根据误差方向设置固定速度
            self.altitude_reached = False
            if error > 0:
                # 需要上升
                cmd_vel.linear.z = self.z_velocity
            else:
                # 需要下降
                cmd_vel.linear.z = -self.z_velocity
                
        # 发布速度命令
        self.cmd_vel_pub.publish(cmd_vel)
        
    def debug_output(self):
        """调试输出"""
        if not self.pose_received:
            self.get_logger().warn('未收到里程计数据，z轴速度为0')
            return
            
        if not self.target_set:
            self.get_logger().info(f'等待目标高度设置... 当前高度: {self.current_altitude:.3f}m')
            return
            
        error = self.target_altitude - self.current_altitude
        z_cmd = 0.0
        if abs(error) > self.altitude_tolerance and self.target_set:
            z_cmd = self.z_velocity if error > 0 else -self.z_velocity
            
        self.get_logger().info(
            f'高度控制 - 当前: {self.current_altitude:.3f}m, '
            f'目标: {self.target_altitude:.3f}m, '
            f'误差: {error:.3f}m, '
            f'z速度: {z_cmd:.1f}m/s, '
            f'状态: {"已到达" if self.altitude_reached else "调整中"}'
        )


def main(args=None):
    rclpy.init(args=args)
    
    controller = SimpleZVelocityController()
    
    try:
        rclpy.spin(controller)
    except KeyboardInterrupt:
        controller.get_logger().info('简单z轴控制器已停止')
    finally:
        controller.destroy_node()
        rclpy.shutdown()


if __name__ == '__main__':
    main() 